from flask_jwt_extended import verify_jwt_in_request, get_jwt_identity, create_access_token, get_jwt
import jwt
from extensions import db
from functools import wraps
from flask import request, current_app, g
from datetime import datetime
from models.token_blacklist import TokenBlacklist

def generate_token(user):
    """使用Flask-JWT生成访问令牌"""
    # 将用户ID转换为字符串作为身份标识
    identity = str(user.user_id)
    additional_claims = {
        'user_id': user.user_id,
        'roles': user.get_roles(),
        'login_type': 'user'
    }
    return str(create_access_token(
        identity=identity,
        additional_claims=additional_claims,
        expires_delta=current_app.config['JWT_ACCESS_TOKEN_EXPIRES']
    ))

def jwt_required():
    """JWT验证装饰器"""
    def wrapper(fn):
        @wraps(fn)
        def decorator(*args, **kwargs):
            verify_jwt_in_request()
            return fn(*args, **kwargs)
        return decorator
    return wrapper

def token_required(f):
    """增强版令牌验证装饰器，支持多类型登录"""
    @wraps(f)
    def decorated(*args, **kwargs):
        from models.user import User
        
        try:
            # 验证并获取JWT数据
            payload = verify_custom_jwt()
            
            # 检查token是否在黑名单中
            jti = payload.get('jti')
            blacklisted = TokenBlacklist.query.filter_by(jti=jti).first()
            if blacklisted:
                return {
                    'status': 'error',
                    'message': '令牌已失效'
                }, 401
                
            # 获取身份信息
            identity = payload.get('user_id', {})  # 使用identity字段获取完整的身份信息
            login_type = payload.get('login_type', 'user')
            
            # 根据登录类型加载身份信息
            if login_type == 'user':
                current_identity = User.query.filter_by(user_id=identity, is_deleted=False).first()
            else:
                return {
                    'status': 'error',
                    'message': '不支持的登录类型'
                }, 403
            
            if not current_identity:
                return {
                    'status': 'error',
                    'message': '身份信息不存在'
                }, 404
            
            # 将登录信息存入请求上下文
            g.current_identity = current_identity
            g.login_type = login_type
            g.user_roles = payload.get('roles', [])
            
            return f(*args, **kwargs)
            
        except jwt.ExpiredSignatureError:
            # 获取过期token的信息并加入黑名单
            try:
                auth_header = request.headers.get('Auuua')
                if auth_header:
                    data = jwt.decode(auth_header, current_app.config['JWT_SECRET_KEY'], algorithms=['HS256'])
                    revoke_token(data['jti'], data['exp'])
            except:
                pass
            return {
                'status': 'error',
                'message': '令牌已过期，请重新登录'
            }, 401
        except jwt.InvalidTokenError as e:
            return {
                'status': 'error',
                'message': str(e)
            }, 401
    
    return decorated

def revoke_token(jti, exp):
    """将令牌加入黑名单"""
    expires = datetime.utcfromtimestamp(exp)
    blacklisted_token = TokenBlacklist(jti=jti, expires=expires)
    db.session.add(blacklisted_token)
    db.session.commit()

def decode_token(token):
    """解码并验证JWT令牌"""
    try:
        payload = jwt.decode(
            token,
            current_app.config['JWT_SECRET_KEY'],
            algorithms=['HS256']
        )
        # 检查令牌是否在黑名单中
        with current_app.app_context():
            if TokenBlacklist.query.filter_by(jti=payload['jti']).first():
                raise jwt.InvalidTokenError('Token revoked')
        return payload
    except jwt.ExpiredSignatureError:
        raise jwt.InvalidTokenError('Token expired')
    except jwt.InvalidTokenError as e:
        raise e
    except Exception as e:
        raise jwt.InvalidTokenError(str(e))

def preprocess_request():
    """请求预处理中间件"""
    print(f"请求路径: {request.path}, 请求方法: {request.method}")
    
    if request.method == 'GET' and request.content_type == 'application/json':
        if not request.data or len(request.data) == 0:
            request._cached_json = (None, False)

def verify_custom_jwt():
    """自定义JWT验证函数"""
    auth_header = (
        request.headers.get('Auuua') or 
        request.headers.get('AUUUA')
    )
    
    if not auth_header:
        raise jwt.InvalidTokenError('未找到认证头')
        
    try:
        # 直接解码验证token
        payload = jwt.decode(
            auth_header,
            current_app.config['JWT_SECRET_KEY'],
            algorithms=['HS256']
        )
        
        # 检查是否在黑名单
        blacklisted = TokenBlacklist.query.filter_by(jti=payload.get('jti')).first()
        if blacklisted:
            raise jwt.InvalidTokenError('令牌已失效')
            
        # 检查是否过期
        exp_timestamp = payload.get('exp')
        if not exp_timestamp:
            raise jwt.InvalidTokenError('令牌缺少过期时间')
            
        # 将解码后的数据存储在g对象中，供后续使用
        g.jwt_claims = payload
        g.jwt_identity = payload.get('sub')
        g.user_roles = payload.get('roles', [])
        
        return payload
        
    except jwt.ExpiredSignatureError:
        raise jwt.ExpiredSignatureError('令牌已过期')
    except jwt.InvalidTokenError as e:
        raise e
    except Exception as e:
        raise jwt.InvalidTokenError(str(e))